#include "../bm/def.h"
#include "lk_api.h"
#include "typedefs.h"

//#define DEBUG_OVER_USB
#define DEBUG_OVER_USB2

#ifdef DEBUG_OVER_USB2
#include "timer.h"
#include "usbdcore.h"
#include "usbtty.h"

static BOOLEAN usb_connect(UINT32 tmo) {
  ULONG start_time = get_timer(0);
  BOOLEAN result = FALSE;
  UINT32 i = 1;

  mt_usb_disconnect_internal();
  mt_usb_connect_internal();

#if CFG_USBIF_COMPLIANCE
  /* USB compliance test: 100mA charging current when USB is unconfigured. */
  platform_set_chrg_cur(70);
#endif

  do {
    /* kick wdt to avoid cpu reset during usb driver installation if not present
     */
    // platform_wdt_all_kick();
    service_interrupts();

    if (usbdl_configured()) {
#if CFG_USBIF_COMPLIANCE
      /* USB compliance test: 500mA charging current when USB is configured but
       * we set the charging current to 450mA since 500mA doesn't support in the
       * platform.
       */
      platform_set_chrg_cur(450);
#endif
      result = TRUE;
      break;
    }

    if (tmo) {
      /* enable timeout mechanism */
      if (get_timer(start_time) > tmo)
        break;
#if !CFG_FPGA_PLATFORM
      /* cable plugged-out and power key detection each 1 second */
      if (get_timer(start_time) > i * 1000) {
        /* Stop booting when usb boot (ex. KPOC) but plug-out usb. */
        // if (!usb_accessory_in() && !pmic_detect_powerkey() && (g_boot_reason
        // == BR_USB))
        //    pl_power_off();

#if !CFG_EVB_PLATFORM
        /* check bypass power key from the 2nd second */
        /*if (i > 1 && pmic_detect_powerkey())
        {
            print("%s Enumeration(Skip): powerkey pressed\n", MOD);
            break;
        }*/
        i++;
#endif
      }
#endif
    }
  } while (1);

  return result;
}
#endif

int (*g_original_part_read)(part_dev_t *dev, uint64_t src, uint8_t *dst,
                            int size, unsigned int part_id);
uint64_t g_boot;
static int part_read_custom(part_dev_t *dev, uint64_t src, uint8_t *dst,
                            int size, unsigned int part_id) {
  video_printf("part_read_custom: 0x%llx 0x%x %d => %p\n", src, size, part_id,
               dst);
  if (src == g_boot) {
    if (size < BLOCK_SIZE * 2)
      return g_original_part_read(dev, src + BLOCK_SIZE * 2, dst, size,
                                  part_id);
    else {
      void *second_copy = (char *)dst + BLOCK_SIZE * 2;
      int r = g_original_part_read(dev, src, second_copy, size, part_id);
      memcpy(dst, second_copy, BLOCK_SIZE * 2);
      memset(second_copy, 0, BLOCK_SIZE * 2);
      return r;
    }
  } else
    return g_original_part_read(dev, src, dst, size, part_id);
}

void patch_lk() {
  volatile uint16_t *x = (volatile uint16_t *)0x4602D3DC;
#if 0
    x[0] = 0xe508; // hook dprint
    arch_clean_invalidate_cache_range((uint32_t)x, 2);

    x = (volatile uint16_t *)0x4600120A;
    x[0] = 0xbf00;
    x[1] = 0xbf00;
    arch_clean_invalidate_cache_range((uint32_t)x, 4);

    x = (volatile uint16_t *)0x4602D9DC; // hook print
    x[0] = 0xf7ff;
    x[1] = 0xba08;
    arch_clean_invalidate_cache_range((uint32_t)x, 8);
#endif

  x = (volatile uint16_t *)0x460222E8; // secure boot disable
  x[0] = 0xf64d;
  x[1] = 0x10b8;
  x[2] = 0xf2c4;
  x[3] = 0x6007;
  x[4] = 0x2100;
  x[5] = 0x6001;
  x[6] = 0xf04f;
  x[7] = 0x0000;
  x[8] = 0x4770;
  arch_clean_invalidate_cache_range((uint32_t)x, 16);
  ((volatile uint32_t *)0x4607D9B8)[0] = 0; // device state: GREEN
  arch_clean_invalidate_cache_range(0x4607D9B8, 4);

  volatile uint32_t *y = (volatile uint32_t *)0x4600120A;
  y[0] = 0xe1a00000;
}

#ifdef DEBUG_OVER_USB
#define usbtty_in (usbtty_eps[0])
#define usbtty_out (usbtty_eps[1])
static struct udc_endpoint *usbtty_eps[2];

static udc_request_t udc_req = {
    .buf = NULL, .length = 0, .complete = NULL, .context = NULL};

volatile int udc_online = 0;
volatile int transfer_done = 0;
volatile int transfer_status = 0;

static void udc_req_complete(udc_request_t *req, unsigned actual, int status) {
  (void)req;

  transfer_status = status;
  req->length = actual;
  transfer_done = 1;

  video_printf("TXD\n");
}

#define MAX_USBFS_BULK_SIZE (16 * 1024)

static int usb_read(void *_buf, uint32_t len) {
  uint8_t *buf = (uint8_t *)_buf;
  int count = 0;

  while (len > 0) {
    uint32_t xfer = (len > MAX_USBFS_BULK_SIZE) ? MAX_USBFS_BULK_SIZE : len;
    udc_req.buf = buf;
    udc_req.length = xfer;
    udc_req.complete = udc_req_complete;
    int status = udc_request_queue(usbtty_out, &udc_req);
    if (status < 0) {
      video_printf("usb_read(): queue failed\n");
      return -1;
    }
    while (transfer_done == 0) {
    }
    transfer_done = 0;

    if (transfer_status < 0) {
      video_printf("usb_read(): transfer failed\n");
      return -1;
    }

    count += udc_req.length;
    buf += udc_req.length;
    len -= udc_req.length;

    if (udc_req.length != xfer)
      break;
  }

  return count;
}

static int usb_write(void *buf, uint32_t len) {
  udc_req.buf = buf;
  udc_req.length = len;
  udc_req.complete = udc_req_complete;
  int status = udc_request_queue(usbtty_in, &udc_req);
  if (status < 0) {
    video_printf("usb_write(): queue failed\n");
    return -1;
  }
  while (transfer_done == 0) {
  }
  transfer_done = 0;

  if (transfer_status < 0) {
    video_printf("usb_write() transfer failed\n");
    return -1;
  }

  return udc_req.length;
}

static void usbtty_notify(struct udc_gadget *gadget, unsigned event) {
  if (event == UDC_EVENT_ONLINE) {
    video_printf("UDC ONLINE\n");
    udc_online = 1;
  } else if (event == UDC_EVENT_OFFLINE) {
    video_printf("UDC OFFLINE\n");
    udc_online = 0;
  }
}

static udc_gadget_t gadget = {
    .notify = usbtty_notify,
    .ifc_class = 0xff,
    .ifc_subclass = 0x42,
    .ifc_protocol = 0x03,
    .ifc_endpoints = 2,
    .ifc_string = "LK Payload debug console",
    .ept = usbtty_eps,
};

void usbtty_init() {
  static udc_device_t dev = {.vendor_id = 0xdee2,
                             .product_id = 0x2eed,
                             .version_id = 0x1000,
                             .manufacturer = "LG",
                             .product = "LK K10 2017 M250E",
                             .serialno = "0123456789"};

  int status;

  video_printf("%s %s %s\n", dev.manufacturer, dev.product, dev.serialno);

  video_printf("UDC init ... ");
  status = udc_init(&dev);
  if (status != 0)
    goto error;
  video_printf("OK\n");

  video_printf("ALLOC eps ... ");
  usbtty_in = udc_endpoint_alloc(UDC_TYPE_BULK_IN, 512);
  if (!usbtty_in)
    goto error_alloc_in;

  usbtty_out = udc_endpoint_alloc(UDC_TYPE_BULK_OUT, 512);
  if (!usbtty_out)
    goto error_alloc_out;
  video_printf("OK\n");

  video_printf("Register gadget ... ");
  status = udc_register_gadget(&gadget);
  if (status != 0)
    goto error;
  video_printf("OK\n");

  video_printf("UDC start ... ");
  status = udc_start();
  if (status != 0)
    goto error;
  video_printf("OK\n");

  video_printf("Wait for UDC online ... ");
  while (udc_online == 0) {
  }
  video_printf("DONE\n");

#if 0
	video_printf("Waiting for tool ... ");
	while (1) {
		char buf[1];
		if (usb_read(buf, sizeof(buf)) == -1) {
			video_printf("ERROR\n");
			for(;;);
		}
		video_printf("?");
		if (buf[0] != 13) {
			video_printf("!");
		} else break;
	}

	video_printf("DONE");
#else
  while (1) {
    char buf[1];
    buf[0] = '#';

    usb_write(buf, sizeof(buf));
  }
#endif

  return;

error:
  video_printf("ERROR status=%d\n", status);
  for (;;)
    ;
error_alloc_in:
  video_printf("ERROR(in)\n");
  for (;;)
    ;
error_alloc_out:
  video_printf("ERROR(out)\n");
  for (;;)
    ;
}
#endif

int main() {
  video_printf("HELLO FROM PAYLOAD\n");

#ifdef DEBUG_OVER_USB
  usbtty_init();
#endif
#ifdef DEBUG_OVER_USB2
  // mtk_timer_init();
  video_printf(".");
  usbdl_init();
  video_printf(".");
  udelay(1000);
  video_printf(".");
  if (!usb_connect(1000)) {
    video_printf("USB INIT ERROR\n");
  } else
    video_printf("USB INIT OK\n");
#endif

#if 1
  volatile part_dev_t *dev = (volatile part_dev_t *)mt_part_get_device();
  if (!dev) {
    video_printf("ERROR\n");
    for (;;)
      ;
  }
  g_original_part_read = dev->read;
  dev->read = part_read_custom;
#endif
  patch_lk();
  video_printf("PATCHING DONE\n");

  void (*mt_boot_init)(void) = (void (*)(void))(0x46024AF4 | 1);
  void (*boot_linux_from_storage)() = (void (*)())(0x46023DCC | 1);

  video_printf("JUMP\n");
  mt_boot_init();
  // boot_linux_from_storage();

  for (;;)
    ;
}
